import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import cv2
import paddlehub
import paddlehub as hub
from PIL import Image
import numpy as np


# print(hub.HubServer())
# paddlehub.server_check()
img_name = 'origin/meditation.jpg' #输入要更换背景的证件照
test_img_path = ["./"+img_name]

module = hub.Module(name="humanseg_mobile", version='1.1.1')


results = module.segment(
    paths=test_img_path,
    visualization=True,
    output_dir='humanseg_output')

print(results)
for result in results:
    print(result)


def blend_images(fore_image, base_image, ):
    # def blend_images(fore_image, base_image):
    """
    将抠出的人物图像换背景
    fore_image: 前景图片，抠出的人物图片
    base_image: 背景图片
    """
    # 读入图片
    base_image = Image.open(base_image).convert('RGB')
    fore_image = Image.open(fore_image).resize(base_image.size)

    # 图片加权合成
    scope_map = np.array(fore_image)[:, :, -1] / 255
    scope_map = scope_map[:, :, np.newaxis]
    scope_map = np.repeat(scope_map, repeats=3, axis=2)
    res_image = np.multiply(scope_map, np.array(fore_image)[:, :, :3]) + np.multiply((1 - scope_map),
                                                                                     np.array(base_image))

    # 保存图片
    res_image = Image.fromarray(np.uint8(res_image))
    res_image.save('./humanseg_output/result.jpg')


blend_images('./humanseg_output/meditation.png', './backgroud/white.png',)